%% EKF估计非线性离散系统
% 由于模型改变，相对于EKF_5_2的区别如下：
% (1) 真值和观测值
% (2) matPhi阵和matHk阵
% (3) 一步预测X(k|k-1)和X估计
clear all;
clc;

%% 参数初始化
N = 200;
T = 0.05;
t = T:T:N*T;
X0 = [1; 0];
P0 = [1, 0; 0, 1];
mu = [0, 0];
matQ = [0.01, 0; 0, 0.0001];
matR = [0.1, 0; 0, 0.1];

%% 生成真值和测量值
rng(1);
Wk = mvnrnd(mu, matQ, N)';
Vk = mvnrnd(mu, matR, N)';
realXk = zeros(2, N);
measureXk = zeros(2, N);

realXk(:, 1) = X0;
measureXk(1, 1) = 2*sin(realXk(1,1)/2) + Vk(1,1);
measureXk(2, 1) = realXk(1,1)/2 + Vk(2,1);
for i = 2:1:N
    x1pre = realXk(1, i-1);
    x2pre = realXk(2, i-1);
    realXk(1, i) = x1pre + T*x2pre + Wk(1, i-1);
    realXk(2, i) = -10*T*sin(x1pre) + (1-T)*x2pre + Wk(2, i-1);
    measureXk(1, i) = 2*sin(realXk(1,i)/2) + Vk(1,i);
    measureXk(2, i) = realXk(1,i)/2 + Vk(2,i);
end

%% EKF
% 初值
estimateXpre = X0;
meanSquareErrorPre = P0;
estimateX = X0;
% 记录
estimateX_EKF = zeros(2, N);
estimateX_EKF(:, 1) = estimateXpre;
estimateMeanSquareError = zeros(2, N);
estimateMeanSquareError(:, 1) = diag(P0);

for i = 2:1:N
    % 相对于KF，EKF的Phi阵要由状态方程f(X)实时求偏导得到
    matPHIk = [1, T; -10*T*cos(estimateXpre(1)), 1-T];
    % 相对于KF，EKF的一部预测方程使用非线性的状态方程f(X)
    estimateXpre = [estimateXpre(1) + T*estimateXpre(2);
                    -10*T*sin(estimateXpre(1)) + (1-T)*estimateXpre(2)];
    % 一步预测Pk
    meanSquareErrorPre = matPHIk*meanSquareErrorPre*matPHIk' + matQ;
    % 相对于KF，EKF的Hk阵要由观测方程h(X)实时求偏导得到
    matHk = [cos(estimateXpre(1)/2), 0; 0.5, 0];
    matHkT = matHk';
    % 增益K
    gainK = meanSquareErrorPre*matHkT/(matHk*meanSquareErrorPre*matHkT + matR);
    % X估计，相对于KF，此处使用观测方程h(X)
    hx = [2*sin(estimateXpre(1)/2); estimateXpre(1)/2];
    estimateX = estimateXpre + gainK*(measureXk(:,i) - hx);
    % P估计
    tmp = eye(2) - gainK*matHk;
    meanSquareErrorK = tmp*meanSquareErrorPre*tmp' + gainK*matR*gainK';
    
    estimateXpre = estimateX;
    meanSquareErrorPre = meanSquareErrorK;
    % 记录
    estimateX_EKF(:, i) = estimateX;
    estimateMeanSquareError(:, i) = diag(meanSquareErrorK);
end

%% 5_4使用matlab自带的EKF
% 初值
estimateXpre = X0;
estimateX_EKF_M = zeros(2, N);
estimateX_EKF_M(:, 1) = estimateXpre;

obj = extendedKalmanFilter(@StateFcn,@MeasurementFcn,estimateX_EKF_M(:,1), ...
        'ProcessNoise', matQ, ...
        'MeasurementNoise', matR, ...
        'StateTransitionJacobianFcn',@StateJacobianFcn, ...
        'MeasurementJacobianFcn',@MeasurementJacobianFcn);

for i = 1:1:N
    % 先correct后prdecit, 记录i而不是i+1
    [CorrectedState, CorrectedStateCovariance] = correct(obj, measureXk(:,i));
    [PredictedState, PredictedStateCovariance] = predict(obj);
    estimateX_EKF_M(:,i) = CorrectedState;
end

%% 画图
% 真值
figure;
plot(t, realXk, 'LineWidth', 1.5);
legend('x1','x2');
title('真值');
% 测量值
figure;
plot(t, measureXk, 'LineWidth', 1.5);
legend('x1','x2');
title('测量值');
% 估计值
figure;
plot(t, estimateX_EKF, 'LineWidth', 1.5);
legend('x1','x2');
title('EKF估计值');
% 估计误差值
figure;
plot(t, estimateX_EKF - realXk, 'LineWidth', 1.5);
legend('x1','x2');
title('估计误差值');
% 真值vs测量值vs估计值
figure;
plot(t, realXk, 'LineWidth', 1.5);
hold on;
plot(t, measureXk, '--', 'LineWidth', 1);
plot(t, estimateX_EKF, 'LineWidth', 1.5);
legend('realx1', 'realx2', 'measurex1', 'measurex2', 'EKFx1', 'EKFx2');
title('真值vs测量值vs估计值');
% 真值vs自写EKFvs官方EKF
figure;
plot(t, realXk, 'LineWidth', 1.5);
hold on;
plot(t, estimateX_EKF, '--', 'LineWidth', 1);
plot(t, estimateX_EKF_M, 'LineWidth', 1.5);
legend('realx1', 'realx2', 'EKFx1', 'EKFx2', 'EKFx1_M', 'EKFx2_M');
title('真值vs自写EKFvs官方EKF');
% 均方差
figure;
plot(t, estimateMeanSquareError, 'LineWidth', 1.5);
legend('x1','x2');
title('均方差');

%% 模型函数
function [ x ] = StateFcn( x )
    T = 0.05;
    x = [x(1)+T*x(2); -10*T*sin(x(1))+(1-T)*x(2)];
end

function [ z ] = MeasurementFcn( x )
    z = [2*sin(x(1)/2); x(1)/2];
end

function [ dfdx ] = StateJacobianFcn( x)
    T = 0.05;    
    dfdx = [1, T; -10*T*cos(x(1)), 1-T];
end

function [ dhdx ] = MeasurementJacobianFcn( x )
    dhdx = [cos(x(1)/2), 0; 1/2, 0];
end
